from pymilvus import connections, Collection, utility
import numpy as np
from typing import List, Dict, Optional

class MilvusClient:
    def __init__(self, host: str = "192.168.2.110", port: str = "19530"):
        """
        初始化 Milvus 客户端
        
        Args:
            host: Milvus 服务器地址
            port: Milvus 服务端口
        """
        self.host = host
        self.port = port
        self.collection = None
        self._connect()
        
    def _connect(self):
        """连接到 Milvus 服务器"""
        try:
            connections.connect(
                alias="default", 
                host=self.host,
                port=self.port
            )
            print(f"成功连接到 Milvus 服务器: {self.host}:{self.port}")
        except Exception as e:
            print(f"连接 Milvus 服务器失败: {str(e)}")
            raise
            
    def load_collection(self, collection_name: str):
        """
        加载指定的集合
        
        Args:
            collection_name: 集合名称
        """
        try:
            if utility.has_collection(collection_name):
                self.collection = Collection(collection_name)
                self.collection.load()
                print(f"成功加载集合: {collection_name}")
            else:
                raise ValueError(f"集合不存在: {collection_name}")
        except Exception as e:
            print(f"加载集合失败: {str(e)}")
            raise
            
    def search(self, 
              vector: np.ndarray,
              limit: int = 5,
              expr: Optional[str] = None) -> List[Dict]:
        """
        执行向量搜索
        
        Args:
            vector: 查询向量
            limit: 返回的最大结果数
            expr: 过滤表达式
            
        Returns:
            搜索结果列表
        """
        if self.collection is None:
            raise ValueError("请先加载集合")
            
        try:
            search_params = {
                "metric_type": "L2",
                "params": {"nprobe": 10}
            }
            
            results = self.collection.search(
                data=[vector],
                anns_field="embedding",
                param=search_params,
                limit=limit,
                expr=expr,
                output_fields=["question", "sql", "ddl"]
            )
            
            search_results = []
            for hits in results:
                for hit in hits:
                    result = {
                        "id": hit.id,
                        "distance": hit.distance,
                        "question": hit.entity.get("question"),
                        "sql": hit.entity.get("sql"),
                        "ddl": hit.entity.get("ddl")
                    }
                    search_results.append(result)
                    
            return search_results
            
        except Exception as e:
            print(f"搜索失败: {str(e)}")
            raise
            
    def close(self):
        """关闭 Milvus 连接"""
        try:
            if self.collection:
                self.collection.release()
            connections.disconnect("default")
            print("已关闭 Milvus 连接")
        except Exception as e:
            print(f"关闭连接失败: {str(e)}")
            
    def __enter__(self):
        """支持 with 语句的上下文管理"""
        return self
        
    def __exit__(self, exc_type, exc_val, exc_tb):
        """退出上下文时自动关闭连接"""
        self.close()



def main():
    """
    主函数，用于测试 Milvus 客户���
    """
    try:
        # 初始化 Milvus 客户端
        client = MilvusClient(
            host="192.168.2.110",
            port="19530"
        )
        
        # 加载集合
        client.load_collection("sql_vectors")
        
        # 执行向量搜索
        query = "查询员工表的所有信息"
        results = client.search(query)
        
        # 打印搜索结果
        print("\n搜索结果:")
        print("-" * 80)
        for i, result in enumerate(results, 1):
            print(f"\n结果 {i}:")
            print(f"相似度距离: {result['distance']}")
            if result.get('question'):
                print(f"问题: {result['question']}")
            if result.get('sql'):
                print(f"SQL: {result['sql']}")
            if result.get('ddl'):
                print(f"DDL: {result['ddl']}")
            print("-" * 80)
            
    except Exception as e:
        print(f"执行失败: {str(e)}")
        import traceback
        traceback.print_exc()
    finally:
        if 'client' in locals():
            client.close()

if __name__ == "__main__":
    main()
